#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/3/7 下午2:00
# @Author  : cherry_wb
# @Site    : 
# @File    : VulDataSet.py
# @Software: PyCharm

from scipy import io
import numpy as np
import os

from dgl.data.dgl_dataset import DGLDataset
from dgl.data.utils import download, save_graphs,save_info, load_graphs, load_info, \
    check_sha1, deprecate_property
from dgl import backend as F
import pickle
import torch
import dgl
import math
from load_bin2graph import load_vul_deepacker,convert_to_one_hot
import dgl


class MDSEQVAEDataSet(DGLDataset):
    _url = 'https://gitee.com/nudt_wb/dataset/mdseqvae_vul.zip'
    _sha1_str = 'xxxx'

    def __init__(self, name='mdseqvae', raw_dir='../dataset/', split=0.0, force_reload=False, verbose=False, max_tolerate_len=1200):
        self.max_tolerate_len = max_tolerate_len
        self.gclasses = 2
        self.nclasses = 1
        self.eclasses = 1
        self.dim_nfeats = 256
        self.split = split

        super(MDSEQVAEDataSet, self).__init__(name=name,
                                         url=self._url,
                                         raw_dir=raw_dir,
                                         force_reload=force_reload,
                                         verbose=verbose)

        if self.split > 0:
            indices = list(range(len(self.graphs)))
            np.random.seed(0)
            np.random.shuffle(indices)
            split_index = int(math.floor(self.split * len(self.graphs)))
            selected_idx = indices[:split_index]
            self.label = self.label[selected_idx]
            self.graphs_new = []
            for indi in indices:
                self.graphs_new.append(self.graphs[indi])

            del self.graphs
            self.graphs = self.graphs_new


    def process(self):
        self.graphs, self.label = self._load_graph()

    def _load_graph(self):
        list_arch_os=['64-ubuntu','32-ubuntu']
        target_dir = "/home/wb/work/paper/paper_code/dataset/mdseqvae/64-ubuntu/cfgedge"
        target_dir_32 = "/home/wb/work/paper/paper_code/dataset/mdseqvae/32-ubuntu/cfgedge"
        list_function_assembly_code, Y_full, max_length, list_function_assembly4onehot = load_vul_deepacker(
            list_arch_os)
        # if max_length < 500:
        #     max_length = 500
        list_function_assembly_onehot = convert_to_one_hot(list_function_assembly4onehot, max_length)
        data_len = len(list_function_assembly_code)
        graphs = []
        labels = []
        label_1_counter = 0
        for i in range(data_len):
            tgt = Y_full[i]
            data = list_function_assembly_onehot[i]
            max_node_id = 0
            for nodedata in data:
                if nodedata.sum() == 0:
                    break
                max_node_id = max_node_id + 1
            edge_index_start = torch.tensor([i for i in range(max_node_id - 1)])
            edge_index_end = torch.tensor([i for i in range(1, max_node_id)])
            edge_index = torch.stack((edge_index_start, edge_index_end))
            edge_attr = torch.tensor([0 for i in range(max_node_id - 1)])

            cfg_file = None
            if i <= 8310 and os.path.exists(os.path.join(target_dir, "{}.cfgedge".format(i))):
                cfg_file = os.path.join(target_dir, "{}.cfgedge".format(i))
            elif i > 8310 and os.path.exists(os.path.join(target_dir_32, "{}.cfgedge".format(i - 8311))):
                cfg_file = os.path.join(target_dir_32, "{}.cfgedge".format(i - 8311))
            else:
                cfg_file = None
            if cfg_file is not None:
                with open(cfg_file, 'r') as f:
                    edges = f.readlines()
                    src = []
                    dst = []
                    for edge in edges:
                        src.append(int(edge.split("->")[0]))
                        dst.append(int(edge.split("->")[1]))
                    # edge_index_cfg = pickle.load(f)
                    edge_index_cfg = torch.stack((torch.tensor(src), torch.tensor(dst)))
                    edge_attr_cfg = torch.tensor([1 for i in range(edge_index_cfg.shape[1])])
                    # edge_attr_cfg = torch.tensor([1 for i in range(edge_index_cfg.shape[1])])
                    edge_index = edge_index_cfg  # torch.cat((edge_index,edge_index_cfg),dim=1)
                    edge_attr = edge_attr_cfg  # torch.cat((edge_attr,edge_attr_cfg))

            g = dgl.graph([])
            g.add_nodes(183)
            g.add_edges(edge_index[0],edge_index[1])

            if tgt > 0:
                label_1_counter = label_1_counter + 1

            g.ndata['attr'] = F.tensor(data, dtype=F.data_type_dict['float32'])
            # g.ndata['annotation'] = F.tensor(nodes_cpg[:src_len])
            g.edata['type'] = F.tensor(edge_attr, dtype=F.data_type_dict['float32'])
            graphs.append(g)
            labels.append(tgt)

        labels = F.tensor(labels)

        return graphs, labels

    def save(self):
        """save the graph list and the labels"""
        graph_path = os.path.join(self.save_path, 'dgl_graph_pyg.bin')
        info_path = os.path.join(self.save_path, 'dgl_graph_pyg.pkl')
        label_dict = {'labels': self.label,}
        info_dict = {
            'gclasses': self.gclasses,
            'nclasses': self.nclasses,
            'eclasses': self.eclasses,
            'dim_nfeats': self.dim_nfeats,
        }
        save_graphs(str(graph_path), self.graphs, label_dict)
        save_info(str(info_path), info_dict)

    def has_cache(self):
        graph_path = os.path.join(self.save_path, 'dgl_graph_pyg.bin')
        return os.path.exists(graph_path)

    def load(self):
        info_dict = load_info(os.path.join(self.save_path, 'dgl_graph_pyg.pkl'))
        graphs, label_dict = load_graphs(os.path.join(self.save_path, 'dgl_graph_pyg.bin'))
        self.graphs = graphs
        self.label = label_dict['labels']
        self.gclasses = info_dict['gclasses']
        self.nclasses = info_dict['nclasses']
        self.eclasses = info_dict['eclasses']
        self.dim_nfeats = info_dict['dim_nfeats']

    def download(self):
        pass

    @property
    def num_labels(self):
        """Number of labels for each graph, i.e. number of prediction tasks."""
        return 2

    def __getitem__(self, idx):
        r""" Get graph and label by index

        Parameters
        ----------
        idx : int
            Item index

        Returns
        -------
        (:class:`dgl.DGLGraph`, Tensor)
        """
        return self.graphs[idx], self.label[idx]

    def __len__(self):
        r"""Number of graphs in the dataset.

        Return
        -------
        int
        """
        return len(self.graphs)


VulData = MDSEQVAEDataSet

if __name__ == '__main__':
    data = VulData()
    print(data.label)
    print(len(data.label))
    print(data.label.sum())